#!/usr/bin/env python3
"""
Python deployment script for web applications
Equivalent to the Node.js version with SSH deployment capabilities
"""

import os
import sys
import zipfile
import shutil
import subprocess
import asyncio
import logging
import re
from datetime import datetime
from pathlib import Path
from typing import Optional, Tuple, List

import paramiko
from dotenv import load_dotenv

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

class DeployError(Exception):
    """部署过程中的自定义异常"""
    pass

class Deployer:
    def __init__(self, git_repo: str, branch_pattern: str, project_name: str = None):
        self.git_repo = git_repo
        self.branch_pattern = branch_pattern
        self.project_name = project_name or self._extract_project_name(git_repo)
        self.project_path = None
        self.ssh_client = None
        self.zip_path = None
        self.backup_path = None
        self.matched_branch = None
        
        # 加载环境变量
        self._load_env()
        
    def _extract_project_name(self, git_repo: str) -> str:
        """从 Git 仓库地址提取项目名称"""
        # 支持多种 Git 仓库格式
        if git_repo.endswith('.git'):
            git_repo = git_repo[:-4]
        
        # 提取最后一部分作为项目名
        project_name = git_repo.split('/')[-1]
        return project_name
    
    def _load_env(self):
        """加载环境变量"""
        env_path = Path(__file__).parent / '.env'
        load_dotenv(env_path)
        
        # 验证必需的环境变量
        required_env_vars = [
            'SERVER_HOST',
            'SERVER_PORT', 
            'SERVER_USERNAME',
            'SERVER_PASSWORD',
            'REMOTE_PATH'
        ]
        
        missing_vars = [var for var in required_env_vars if not os.getenv(var)]
        if missing_vars:
            raise DeployError(f"缺少必需的环境变量: {', '.join(missing_vars)}")
            
        self.config = {
            'host': os.getenv('SERVER_HOST'),
            'port': int(os.getenv('SERVER_PORT', '22')),
            'username': os.getenv('SERVER_USERNAME'),
            'password': os.getenv('SERVER_PASSWORD'),
            'remote_path': os.getenv('REMOTE_PATH'),
            'backup_path': os.getenv('BACKUP_PATH') or f"{os.getenv('REMOTE_PATH')}/backups"
        }
    
    def log(self, message: str, level: str = 'info'):
        """日志记录函数"""
        emoji_map = {
            'error': '❌',
            'success': '✅', 
            'warning': '⚠️',
            'info': 'ℹ️'
        }
        emoji = emoji_map.get(level, 'ℹ️')
        logger.info(f"{emoji} {message}")
    
    async def sleep(self, seconds: float):
        """异步睡眠"""
        await asyncio.sleep(seconds)
    
    def exec_command(self, command: str, cwd: Optional[str] = None) -> str:
        """执行本地命令"""
        try:
            result = subprocess.run(
                command,
                shell=True,
                cwd=cwd or str(self.project_path),
                capture_output=True,
                text=True,
                encoding='utf-8',
                errors='replace',
                check=True
            )
            return result.stdout.strip() if result.stdout else ""
        except subprocess.CalledProcessError as e:
            stderr = e.stderr.strip() if e.stderr else "未知错误"
            raise DeployError(f"命令执行失败: {command}\n错误: {stderr}")
    
    def get_available_branches(self) -> List[str]:
        """获取远程仓库的所有分支"""
        try:
            # 获取远程分支信息
            result = subprocess.run(
                f'git ls-remote --heads {self.git_repo}',
                shell=True,
                capture_output=True,
                text=True,
                encoding='utf-8',
                errors='replace',
                check=True
            )
            
            branches = []
            stdout = result.stdout.strip() if result.stdout else ""
            for line in stdout.split('\n'):
                if line:
                    # 提取分支名（去掉 refs/heads/ 前缀）
                    branch = line.split('\t')[-1].replace('refs/heads/', '')
                    branches.append(branch)
            
            return branches
            
        except subprocess.CalledProcessError as e:
            stderr = e.stderr.strip() if e.stderr else "未知错误"
            raise DeployError(f"获取分支列表失败: {stderr}")
    
    def find_matching_branch(self, branches: List[str]) -> str:
        """根据模式查找匹配的分支"""
        matching_branches = []
        
        # 将模式转换为正则表达式
        pattern = re.compile(self.branch_pattern, re.IGNORECASE)
        
        for branch in branches:
            if pattern.search(branch):
                matching_branches.append(branch)
        
        if not matching_branches:
            available_branches = ', '.join(branches[:10])  # 只显示前10个分支
            if len(branches) > 10:
                available_branches += f' ... (共 {len(branches)} 个分支)'
            raise DeployError(f"没有找到匹配模式 '{self.branch_pattern}' 的分支。可用分支: {available_branches}")
        
        if len(matching_branches) > 1:
            self.log(f"找到多个匹配的分支: {', '.join(matching_branches)}", 'warning')
            # 优先选择 main 或 master 分支
            for priority_branch in ['main', 'master']:
                if priority_branch in matching_branches:
                    self.log(f"选择优先级分支: {priority_branch}", 'info')
                    return priority_branch
            
            # 如果没有优先级分支，选择第一个
            selected_branch = matching_branches[0]
            self.log(f"选择第一个匹配的分支: {selected_branch}", 'info')
            return selected_branch
        
        return matching_branches[0]
    
    async def clone_or_update_repo(self):
        """克隆或更新 Git 仓库"""
        try:
            # 检查 Git 是否可用
            try:
                self.exec_command('git --version', str(Path.cwd()))
                self.log('Git 可用', 'info')
            except Exception:
                raise DeployError("Git 不可用，请确保已安装 Git")
            
            # 创建临时工作目录
            temp_dir = Path.cwd() / 'temp_deploy'
            temp_dir.mkdir(exist_ok=True)
            
            project_dir = temp_dir / self.project_name
            
            if project_dir.exists():
                self.log(f"项目目录已存在，更新代码...")
                self.project_path = project_dir
                
                # 获取可用分支
                branches = self.get_available_branches()
                self.log(f"获取到 {len(branches)} 个远程分支")
                
                # 查找匹配的分支
                self.matched_branch = self.find_matching_branch(branches)
                self.log(f"选择分支: {self.matched_branch}", 'success')
                
                # 切换到指定分支
                self.exec_command(f'git checkout {self.matched_branch}')
                self.log(f'切换到分支: {self.matched_branch}', 'success')
                
                # 拉取最新代码
                self.exec_command('git pull origin ' + self.matched_branch)
                self.log('代码更新成功', 'success')
                
            else:
                self.log(f"克隆仓库: {self.git_repo}")
                
                # 获取可用分支
                branches = self.get_available_branches()
                self.log(f"获取到 {len(branches)} 个远程分支")
                
                # 查找匹配的分支
                self.matched_branch = self.find_matching_branch(branches)
                self.log(f"选择分支: {self.matched_branch}", 'success')
                
                # 克隆指定分支
                self.exec_command(f'git clone -b {self.matched_branch} {self.git_repo} {self.project_name}', str(temp_dir))
                self.project_path = project_dir
                self.log('仓库克隆成功', 'success')
            
        except Exception as e:
            self.log(f"Git 操作失败: {str(e)}", 'error')
            raise
    
    async def build_project(self):
        """构建项目"""
        try:
            # 检查项目路径是否存在
            if not self.project_path.exists():
                raise DeployError(f"项目路径 {self.project_path} 不存在!")
            
            # 检查 pnpm 是否可用
            try:
                self.exec_command('pnpm --version')
                self.log('pnpm 可用', 'info')
            except Exception:
                raise DeployError("pnpm 不可用，请确保已安装 pnpm")
            
            # 检查 package.json 是否存在
            package_json = self.project_path / 'package.json'
            if not package_json.exists():
                raise DeployError("项目目录中未找到 package.json 文件")
            
            self.log('开始安装依赖...')
            self.log('执行安装命令: pnpm install')
            
            # 先安装依赖
            self.exec_command('npm install')
            self.log('依赖安装成功!', 'success')
            
            self.log('开始构建项目...')
            self.log('执行构建命令: pnpm build')
            
            # 执行构建
            self.exec_command('npm run build:prod')
            self.log('构建成功!', 'success')
            
        except Exception as e:
            self.log(f"构建失败: {str(e)}", 'error')
            raise
    
    def create_zip(self) -> str:
        """创建压缩包"""
        dist_path = self.project_path / 'dist'
        self.zip_path = self.project_path / 'dist.zip'
        
        self.log(f"正在压缩: {dist_path} -> {self.zip_path}")
        
        # 检查dist目录是否存在
        if not dist_path.exists():
            raise DeployError(f"构建输出目录 {dist_path} 不存在!")
        
        try:
            with zipfile.ZipFile(self.zip_path, 'w', zipfile.ZIP_DEFLATED, compresslevel=9) as zipf:
                for file_path in dist_path.rglob('*'):
                    if file_path.is_file():
                        arcname = file_path.relative_to(dist_path)
                        zipf.write(file_path, arcname)
            
            size_mb = self.zip_path.stat().st_size / (1024 * 1024)
            self.log(f"压缩完成! 文件大小: {size_mb:.2f}MB", 'success')
            return str(self.zip_path)
            
        except Exception as e:
            raise DeployError(f"压缩失败: {str(e)}")
    
    async def connect_ssh(self):
        """连接SSH服务器"""
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            self.log('正在连接服务器...')
            self.ssh_client.connect(
                hostname=self.config['host'],
                port=self.config['port'],
                username=self.config['username'],
                password=self.config['password'],
                timeout=30
            )
            self.log('服务器连接成功!', 'success')
            
        except Exception as e:
            raise DeployError(f"SSH连接失败: {str(e)}")
    
    def ssh_exec(self, command: str) -> Tuple[int, str, str]:
        """执行SSH命令"""
        if not self.ssh_client:
            raise DeployError("SSH连接未建立")
        
        stdin, stdout, stderr = self.ssh_client.exec_command(command)
        exit_code = stdout.channel.recv_exit_status()
        
        # 安全地读取输出，处理编码问题
        try:
            stdout_text = stdout.read().decode('utf-8', errors='replace').strip()
        except Exception:
            stdout_text = ""
            
        try:
            stderr_text = stderr.read().decode('utf-8', errors='replace').strip()
        except Exception:
            stderr_text = ""
            
        return exit_code, stdout_text, stderr_text
    
    async def backup_server(self) -> Optional[str]:
        """备份服务器当前版本"""
        try:
            self.log('开始备份服务器当前版本...')
            
            # 创建备份目录
            backup_dir = self.config['backup_path']
            self.ssh_exec(f"mkdir -p {backup_dir}")
            
            # 生成备份文件名（包含时间戳）
            timestamp = datetime.now().isoformat().replace(':', '-').replace('.', '-')
            backup_name = f"backup-{timestamp}.zip"
            backup_path = f"{backup_dir}/{backup_name}"
            
            # 检查当前部署目录是否存在，如果存在则备份
            exit_code, stdout, stderr = self.ssh_exec(
                f"test -d {self.config['remote_path']} && echo 'exists' || echo 'not_exists'"
            )
            
            if stdout == 'exists':
                # 备份dist目录（进入dist目录，压缩其内容）
                self.ssh_exec(f"cd {self.config['remote_path']} && zip -r {backup_path} .")
                self.log(f"服务器备份完成: {backup_name}", 'success')
                
                # 清理旧的备份文件（保留最近5个）
                self.ssh_exec(f"cd {backup_dir} && ls -t | tail -n +6 | xargs -r rm -f")
                self.log('清理旧备份文件完成', 'success')
                
                return backup_path
            else:
                self.log('当前服务器没有可备份的文件', 'warning')
                return None
                
        except Exception as e:
            self.log(f"备份失败: {str(e)}", 'error')
            # 备份失败不应该阻止部署继续
            return None
    
    async def rollback_server(self, backup_path: str):
        """回滚服务器"""
        if not backup_path:
            self.log('没有可用的备份文件进行回滚', 'warning')
            return
        
        try:
            self.log('开始回滚到上一个版本...')
            
            # 回滚到备份版本
            self.ssh_exec(f"""
                cd {self.config['remote_path']} && \
                cd .. && \
                rm -rf dist && \
                mkdir -p dist && \
                cd dist && \
                unzip -o {backup_path} && \
                echo "回滚完成"
            """)
            
            self.log('回滚成功!', 'success')
            
        except Exception as e:
            self.log(f"回滚失败: {str(e)}", 'error')
    
    def upload_file(self, local_path: str, remote_path: str):
        """上传文件到服务器"""
        try:
            sftp = self.ssh_client.open_sftp()
            sftp.put(local_path, remote_path)
            sftp.close()
            self.log('文件上传成功!', 'success')
        except Exception as e:
            raise DeployError(f"文件上传失败: {str(e)}")
    
    async def deploy(self):
        """执行部署"""
        try:
            # 克隆或更新仓库
            await self.clone_or_update_repo()
            
            # 构建项目
            await self.build_project()
            
            # 创建压缩包
            zip_path = self.create_zip()
            
            # 连接服务器
            await self.connect_ssh()
            
            # 备份当前版本
            self.backup_path = await self.backup_server()
            
            # 上传文件
            self.log('开始上传文件...')
            self.upload_file(zip_path, f"{self.config['remote_path']}/dist.zip")
            
            # 部署新版本
            self.log('正在部署新版本...')
            exit_code, stdout, stderr = self.ssh_exec(f"""
                cd {self.config['remote_path']} && \
                rm -rf dist && \
                unzip -o dist.zip && \
                rm dist.zip && \
                echo "部署完成"
            """)
            
            if exit_code == 0:
                self.log('部署完成!', 'success')
                self.log(f"部署输出: {stdout}")
            else:
                raise DeployError(f"部署命令执行失败: {stderr}")
                
        except Exception as e:
            self.log(f"部署过程中出现错误: {str(e)}", 'error')
            
            # 如果部署失败且有备份，尝试回滚
            if self.backup_path:
                self.log('尝试回滚到上一个版本...', 'warning')
                try:
                    await self.rollback_server(self.backup_path)
                except Exception as rollback_error:
                    self.log(f"回滚也失败了: {str(rollback_error)}", 'error')
            
            raise
            
        finally:
            # 清理本地临时文件
            if self.zip_path and self.zip_path.exists():
                self.zip_path.unlink()
                self.log('本地临时文件已清理', 'success')
            
            # 关闭SSH连接
            if self.ssh_client:
                self.ssh_client.close()

def main():
    """主函数"""
    # 检查命令行参数
    if len(sys.argv) < 3:
        print('请指定 Git 仓库地址和分支模式!')
        print('使用方法: python main.py <git仓库地址> <分支模式> [项目名称]')
        print('示例: python main.py https://github.com/user/repo.git "feature/*"')
        print('示例: python main.py https://github.com/user/repo.git "main" my-project')
        sys.exit(1)
    
    git_repo = sys.argv[1]
    branch_pattern = sys.argv[2]
    project_name = sys.argv[3] if len(sys.argv) > 3 else None
    
    # 显示调试信息
    print(f"Git 仓库: {git_repo}")
    print(f"分支模式: {branch_pattern}")
    print(f"项目名称: {project_name}")
    print(f"当前工作目录: {os.getcwd()}")
    
    try:
        # 创建部署器实例
        deployer = Deployer(git_repo, branch_pattern, project_name)
        
        # 执行部署
        asyncio.run(deployer.deploy())
        
        deployer.log('所有操作完成!', 'success')
        sys.exit(0)
        
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
        # 打印详细的错误信息
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()
